from python_a2a.server import A2AServer
from python_a2a.models import Message, TextContent, MessageRole
from python_a2a import run_server
import re
import requests

class DuckDuckGoAgent(A2AServer):
    """Agent that finds stock ticker symbols."""
    
    def __init__(self):
        super().__init__(url="http://localhost:5003")
        self.mcp_server_url = "http://localhost:5001"
    
    def handle_message(self, message):
        """同步处理消息"""
        if hasattr(message.content, 'text'):
            text = message.content.text
            
            # 提取公司名称
            company_name = self._extract_company_name(text)
            
            if not company_name:
                return Message(
                    content=TextContent(text="请提供公司名称来查找股票代码。"),
                    role=MessageRole.AGENT,
                    parent_message_id=getattr(message, 'message_id', None),
                    conversation_id=getattr(message, 'conversation_id', None)
                )
            
            # 使用MCP工具查找股票代码
            ticker = self._call_mcp_tool(company_name)
            
            return Message(
                content=TextContent(text=f"{ticker['content'][0]['text']}"),
                role=MessageRole.AGENT,
                parent_message_id=getattr(message, 'message_id', None),
                conversation_id=getattr(message, 'conversation_id', None)
            )
        
        
        return Message(
            content=TextContent(text="我可以帮助查找公司的股票代码。"),
            role=MessageRole.AGENT,
            parent_message_id=getattr(message, 'message_id', None),
            conversation_id=getattr(message, 'conversation_id', None)
        )
    
    def _call_mcp_tool(self, company_name: str) -> str:
        """直接调用MCP服务器查找股票代码"""
        payload = {
            "jsonrpc": "2.0",
            "id": 1,
            "method": "tools/call",
            "params": {
                "name": "search_ticker",
                "arguments": {
                    "company_name": company_name
                }
            }
        }
        
        response = requests.post(
            self.mcp_server_url,
            json=payload,
            headers={"Content-Type": "application/json"},
            timeout=10
        )
        
        if response.status_code == 200:
            result = response.json()
            
            if "result" in result:
                ticker = result["result"]
                return ticker
            elif "error" in result:
                error_msg = f"MCP服务器错误: {result['error']}"
                raise Exception(error_msg)
            
    
    def _extract_company_name(self, text: str) -> str:
        """从文本中提取公司名称"""
        patterns = [
            r"(.+?)的股票代码",
            r"股票代码(.+?)",
            r"(.+?)的股票",
            r"(.+?)股票代码",
            r"ticker\s+(?:for|of)\s+(.+)",
            r"stock\s+(?:symbol|ticker)\s+(?:for|of)\s+(.+)",
            r"symbol\s+for\s+(.+)"
        ]
        
        text_lower = text.lower()
        
        for pattern in patterns:
            match = re.search(pattern, text_lower, re.IGNORECASE)
            if match:
                company = match.group(1).strip()
                company = re.sub(r'[?？。.,!！]', '', company)
                return company
        
        cleaned_text = re.sub(r'(什么|多少|哪个|哪只|查找|查询|搜索|的股票代码|的股票|股票代码|ticker|stock|symbol)', '', text_lower, flags=re.IGNORECASE)
        cleaned_text = cleaned_text.strip()
        
        return cleaned_text if cleaned_text else ""

if __name__ == "__main__":
    run_server(DuckDuckGoAgent(), port=5003)